#include "stdafx.h"

#include "HookWinsock.h"
#include "HookLoadLibrary.h"

LoadLibraryA_t   pTrueLoadLibraryA;
LoadLibraryW_t   pTrueLoadLibraryW;
LoadLibraryExA_t pTrueLoadLibraryExA;
LoadLibraryExW_t pTrueLoadLibraryExW;

extern LPCSTR  DLL_TOHOOK;
extern LPCWSTR LDLL_TOHOOK;

void UnhookLoadLibrary();

/////////////////////////////////////////////////////////////////////
HMODULE WINAPI HLoadLibraryA(char* str)
{
	HMODULE hMod = pTrueLoadLibraryA(str);
	if(!strcmp(str,DLL_TOHOOK))
	{
		GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_PIN,DLL_TOHOOK,&hMod);
		try
        {
            HookWinsock2Dll(hMod);
        }
        catch (std::exception& ex)
        {
            LOG(plog::error) << ex.what();
        }
		UnhookLoadLibrary();
	}
    return hMod;
}

/////////////////////////////////////////////////////////////////////
HMODULE WINAPI HLoadLibraryW(wchar_t* wstr)
{
	HMODULE hMod=pTrueLoadLibraryW(wstr);
	if(!wcscmp(wstr,LDLL_TOHOOK))
	{
		GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_PIN,LDLL_TOHOOK,&hMod);
		try
        {
            HookWinsock2Dll(hMod);
        }
        catch (std::exception& ex)
        {
            LOG(plog::error) << ex.what();
        }
		UnhookLoadLibrary();
	}
    return hMod;
}

/////////////////////////////////////////////////////////////////////
HMODULE WINAPI HLoadLibraryExA(char* str, HANDLE x, DWORD dwFlags)
{
	HMODULE hMod=pTrueLoadLibraryExA(str,NULL,dwFlags);
	if(!strcmp(str,DLL_TOHOOK))
	{
		GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_PIN,DLL_TOHOOK,&hMod);
		try
        {
            HookWinsock2Dll(hMod);
        }
        catch (std::exception& ex)
        {
            LOG(plog::error) << ex.what();
        }
		UnhookLoadLibrary();
	}
    return hMod;
}

/////////////////////////////////////////////////////////////////////
HMODULE WINAPI HLoadLibraryExW(wchar_t* wstr, HANDLE x, DWORD dwFlags)
{
	HMODULE hMod=pTrueLoadLibraryExW(wstr,x,dwFlags);
	if(!wcscmp(wstr,LDLL_TOHOOK))
	{
		GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_PIN,LDLL_TOHOOK,&hMod);
		try
        {
            HookWinsock2Dll(hMod);
        }
        catch (std::exception& ex)
        {
            LOG(plog::error) << ex.what();
        }
		UnhookLoadLibrary();
	}
    return hMod;
}

/////////////////////////////////////////////////////////////////////
void HookLoadLibrary ()
{
    pTrueLoadLibraryA =
        (LoadLibraryA_t)GetProcAddress(GetModuleHandle(L"Kernel32.dll"),"LoadLibraryA");
    pTrueLoadLibraryW =
        (LoadLibraryW_t)GetProcAddress(GetModuleHandle(L"Kernel32.dll"),"LoadLibraryW");
    pTrueLoadLibraryExA =
        (LoadLibraryExA_t)GetProcAddress(GetModuleHandle(L"Kernel32.dll"),"LoadLibraryExA");
    pTrueLoadLibraryExW =
        (LoadLibraryExW_t)GetProcAddress(GetModuleHandle(L"Kernel32.dll"),"LoadLibraryExW");
    
    if (!Mhook_SetHook((PVOID*)&pTrueLoadLibraryA,HLoadLibraryA))
    {
        throw std::runtime_error ("Mhook HLoadLibraryA failure");
    }
    if (!Mhook_SetHook((PVOID*)&pTrueLoadLibraryW,HLoadLibraryW))
    {
        throw std::runtime_error ("Mhook HLoadLibraryW failure");
    }
    if (!Mhook_SetHook((PVOID*)&pTrueLoadLibraryExA,HLoadLibraryExA))
    {
        throw std::runtime_error ("Mhook HLoadLibraryExA failure");
    }
    if (!Mhook_SetHook((PVOID*)&pTrueLoadLibraryExW,HLoadLibraryExW))
    {
        throw std::runtime_error ("Mhook HLoadLibraryExW failure");
    }
}

/////////////////////////////////////////////////////////////////////
void UnhookLoadLibrary()
{
	Mhook_Unhook((PVOID*)&pTrueLoadLibraryA);
	Mhook_Unhook((PVOID*)&pTrueLoadLibraryW);
	Mhook_Unhook((PVOID*)&pTrueLoadLibraryExA);
	Mhook_Unhook((PVOID*)&pTrueLoadLibraryExW);
}
